import numpy as np  
import cv2 
import time 
import matplotlib.pyplot as plt
import scipy.signal as signal


from time import clock  
 # 为使用的图像金字塔层数 
lk_params = dict( winSize  = (15, 15),  #窗口大小 
                  maxLevel = 4,   #金字塔层数 可以有效面对快速运动的图片
                  criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))      
  
feature_params = dict( maxCorners = 100,   
                       qualityLevel = 0.3,  
                       minDistance = 7,  
                       blockSize = 7 )  
  
class App:  
    def __init__(self, video_src):#构造方法，初始化一些参数和视频路径  
        self.track_len = 10  
        self.detect_interval = 1  
        self.tracks = []  
        self.cam = cv2.VideoCapture(video_src)  
        self.frame_idx = 0
        self.speed1x=[]
        self.speed1y=[]
        self.speed2x=[]
        self.speed2y=[]
        self.wy=[]
        self.wy_y=[]
        self.new_tracks = []
        self.flag=0 
        self.color=0 
        self.hx=[]
        self.hy=[]
        
    def run(self):#光流运行方法

        while True:  
            ret, frame = self.cam.read()#读取视频帧 
            print(ret)
            if ret == True:  
                
                frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)#转化为灰度虚图像 

                vis = frame.copy()  
      
                if len(self.tracks) > 0:#检测到角点后进行光流跟踪  
                    img0, img1 = self.prev_gray, frame_gray  
                    p0 = np.float32([tr[-1] for tr in self.tracks]).reshape(-1, 1, 2)#p0 一直在增加，每5帧加50 
                    p1, st, err = cv2.calcOpticalFlowPyrLK(img0, img1, p0, None, **lk_params)
                    t_start = time.time()
                    M, mask = cv2.findHomography(p0, p1, cv2.RANSAC,5.0)
                    good=mask
                    self.new_tracks = []#数组会被清零
                    for tr, (x, y), good_flag in zip(self.tracks, p1.reshape(-1, 2), good):#将跟踪正确的点列入成功跟踪点 
                        #if not good_flag:  #注释 不滤波
                        #    continue 
                        tr.append((x, y))  
                        if len(tr) > self.track_len:  
                        #    pass
                            del tr[0]##好像会循环5次 每次更新5个点
                            #del tr[1]
                        self.new_tracks.append(tr)
                        cv2.circle(vis, (x, y), 2, (0, 255, 0), -1) 
                    cv2.polylines(vis, [np.int32(tr) for tr in self.tracks], False, (0, 0, 255),5)#以上一振角点为初始点，当前帧跟踪到的点为终点划线  
                if 1:#每5帧检测一次特征点  
                    mask = np.zeros_like(frame_gray)#初始化和视频大小相同的图像  
                    mask[:] = 255#将mask赋值255也就是算全部图像的角点  
                    for x, y in [np.int32(tr[-1]) for tr in self.tracks]:#跟踪的角点画圆  
                        cv2.circle(mask, (x, y), 5, 0, -1)
                    '''每5帧清空位移保存数组，每5帧计算一次位移 '''
                    self.flag+=1
                    self.speed1x=[]
                    self.speed2x=[]
                    self.speed1y=[]
                    self.speed2y=[]
                    for i in range(0, len(self.new_tracks)):
                        self.speed1x.append(self.new_tracks[i][0][0])
                        self.speed1y.append(self.new_tracks[i][0][1])
                        self.speed2x.append(self.new_tracks[i][-1][0])
                        self.speed2y.append(self.new_tracks[i][-1][1])
                    speed1x = np.array(np.array(self.speed1x, dtype='float32'))
                    speed1y = np.array(np.array(self.speed1y, dtype='float32'))
                    speed2x = np.array(np.array(self.speed2x, dtype='float32'))
                    speed2y = np.array(np.array(self.speed2y, dtype='float32'))
                    Xi=speed2x-speed1x
                    Yi=speed2y-speed1y
                    self.wy.append( np.nan_to_num( np.mean(Xi) )  )#nan 转换成数字
                    self.wy_y.append( np.nan_to_num( np.mean(Yi) )  )#nan 转换成数字
                    a=np.array( sum(self.wy[1:]) ,dtype='float16')
                    b=np.array( sum(self.wy_y[1:]) ,dtype='float16')
                    self.hx.append(a)
                    self.hy.append(b)
                    if self.flag%(10)==0:
                        plt.figure("tu1") 
                        #plt.plot(self.hx, self.hy,marker=maker[self.color],c=color[self.color])#scatter
                        #plt.plot(self.hx, self.hy,c=color[self.color])
                        plt.plot(self.hx, self.hy,c="r")
                        plt.draw()
                        plt.pause(0.01)
                        #print(self.hx,self.hy)
                        self.hx=[]
                        self.hy=[]
                        self.color+=1
                    self.tracks = []
                    
                    p = cv2.goodFeaturesToTrack(frame_gray, mask = mask, **feature_params)#像素级别角点检测  
                    if p is not None:  
                        for x, y in np.float32(p).reshape(-1, 2):  
                            self.tracks.append([(x, y)])#将检测到的角点放在待跟踪序列中  
                    

                self.frame_idx += 1  
                self.prev_gray = frame_gray  
                cv2.imshow('lk_track', vis)  
            ch = 0xFF & cv2.waitKey(1)  #按键q
            if ch == 27: 
                break  
def main():  
    import sys  
    #video_src = "E:\Code_Py\Video16.MP4"
    #video_src = "E:\\002.wmv"
    video_src = 0
    #print (__doc__) 
    global orb,fast,freakExtractor
    global color,maker
    color=['b','g','r','c','m','y','k','w','b','g']
    maker=['$A$','$B$','$C$','$M$','$R$','$y$','$k$','$Q$','$W$','$E$']
    App(video_src).run()  
    cv2.destroyAllWindows()               

if __name__ == '__main__':  
    main()  